package com.supreme.secretserver.component;

import cn.hutool.core.io.IoUtil;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.supreme.common.constant.ConsoleColors;
import com.supreme.secretserver.annotation.Decrypt;
import com.supreme.secretserver.props.EncryptProps;
import com.supreme.secretserver.secret.SecretParser;
import org.apache.commons.lang3.StringUtils;
import org.springframework.beans.factory.annotation.Qualifier;
import org.springframework.core.MethodParameter;
import org.springframework.http.HttpInputMessage;
import org.springframework.http.converter.HttpMessageConverter;
import org.springframework.http.converter.json.MappingJacksonInputMessage;
import org.springframework.web.bind.annotation.ControllerAdvice;
import org.springframework.web.servlet.mvc.method.annotation.RequestBodyAdvice;

import javax.annotation.Resource;
import java.io.IOException;
import java.io.InputStream;
import java.lang.reflect.Type;
import java.util.Objects;
import java.util.Optional;

@ControllerAdvice
public class DecryptRequest implements RequestBodyAdvice {

    private final ObjectMapper objectMapper = new ObjectMapper();

    @Resource
    private EncryptProps encryptProps;

    @Resource
    @Qualifier(value = "aesSecretParser")
    private SecretParser aesSecretParser;

    @Resource
    @Qualifier(value = "rsaSecretParser")
    private SecretParser rsaSecretParser;

    @Override
    public boolean supports(MethodParameter methodParameter, Type type, Class<? extends HttpMessageConverter<?>> clazz) {
        if (methodParameter.hasMethodAnnotation(Decrypt.class)) {
            Decrypt encryptAnno = methodParameter.getMethodAnnotation(Decrypt.class);
            Boolean decryptEnable = Optional.ofNullable(encryptAnno)
                    .filter(Decrypt::enable)
                    .filter(encrypt -> StringUtils.equalsAny(encrypt.type(), SecretParser.AES, SecretParser.RSA))
                    .isPresent();
            ConsoleColors.println(ConsoleColors.CYAN_BOLD_BRIGHT, "EncryptResponse decryptEnable: " + decryptEnable);
            return decryptEnable;
        }

        return false;
    }

    @Override
    public HttpInputMessage beforeBodyRead(HttpInputMessage httpInputMessage, MethodParameter methodParameter, Type type, Class<? extends HttpMessageConverter<?>> clazz) throws IOException {
        InputStream inputMessageBody = httpInputMessage.getBody();

        Decrypt methodAnnotation = methodParameter.getMethodAnnotation(Decrypt.class);
        if (Optional.ofNullable(methodAnnotation).isPresent()) {
            try {
                String secretType = methodAnnotation.type();
                String decryptKey = methodAnnotation.decryptKey();

                byte[] bodyBytes = IoUtil.readBytes(inputMessageBody);
                String bodyContent = new String(bodyBytes);

                byte[] decryptKeyBytes = new byte[0];
                SecretParser secretParser = null;

                if (SecretParser.AES.equals(secretType)) {
                    String secretKey = methodAnnotation.secretKey();
                    decryptKeyBytes = Optional.ofNullable(secretKey)
                            .filter(StringUtils::isNotBlank)
                            .map(String::getBytes)
                            .orElse(encryptProps.getAesKeyBytes());

                    secretParser = aesSecretParser;
                } else if (SecretParser.RSA.equals(secretType)) {
                    String secretKey = methodAnnotation.privateKey();
                    decryptKeyBytes = Optional.ofNullable(secretKey)
                            .filter(StringUtils::isNotBlank)
                            .map(String::getBytes)
                            .orElse(encryptProps.getRsaPriKeyBytes());

                    secretParser = rsaSecretParser;
                }

                if (Objects.nonNull(secretParser)) {
                    if (SecretParser.isJsonBody(bodyContent)) {
                        InputStream inputStream = secretParser.buildJsonInputStream(bodyBytes, objectMapper, decryptKey, decryptKeyBytes);
                        HttpInputMessage decryHttpMessage = new MappingJacksonInputMessage(inputStream, httpInputMessage.getHeaders());
                        return decryHttpMessage;
                    } else {
                        InputStream inputStream = secretParser.buildStringInputStream(bodyContent, decryptKey, decryptKeyBytes);
                        HttpInputMessage decryHttpMessage = new MappingJacksonInputMessage(inputStream, httpInputMessage.getHeaders());
                        return decryHttpMessage;
                    }
                }
            } catch (Exception e) {
                e.printStackTrace();
            }
        }

        return httpInputMessage;
    }

    @Override
    public Object afterBodyRead(Object o, HttpInputMessage httpInputMessage, MethodParameter methodParameter, Type type, Class<? extends HttpMessageConverter<?>> aClass) {
        return o;
    }

    @Override
    public Object handleEmptyBody(Object o, HttpInputMessage httpInputMessage, MethodParameter methodParameter, Type type, Class<? extends HttpMessageConverter<?>> aClass) {
        return o;
    }
}
